// arch/riscv/kernel/vm.c

#include "vm.h"
#include "defs.h"
#include "string.h"
#include "mm.h"
#include "printk.h"

extern char _stext[];
extern char _srodata[];
extern char _sdata[];
extern char _sbss[];


/* early_pgtbl: 用于 setup_vm 进行 1GB 的 映射。 */
unsigned long early_pgtbl[512] __attribute__((__aligned__(0x1000)));
/* swapper_pg_dir: kernel pagetable 根目录， 在 setup_vm_final 进行映射。 */
unsigned long swapper_pg_dir[512] __attribute__((__aligned__(0x1000)));

void setup_vm(void)
{
    /*
    1. 由于是进行 1GB 的映射 这里不需要使用多级页表
    2. 将 va 的 64bit 作为如下划分： | high bit | 9 bit | 30 bit |
        high bit 可以忽略
        中间9 bit 作为 early_pgtbl 的 index
        低 30 bit 作为 页内偏移 这里注意到 30 = 9 + 9 + 12， 即我们只使用根页表， 根页表的每个 entry 都对应 1GB 的区域。
    3. Page Table Entry 的权限 V | R | W | X 位设置为 1
    */
    uint64 index,ppn2;
    ppn2 = (uint64)(PHY_START & 0xffffffc0000000) >> 2;
    // 等值映射
    index = ((uint64)PHY_START >> 30) & 511;
    early_pgtbl[index] = ppn2 | PTE_V | PTE_R | PTE_W | PTE_X | PTE_A | PTE_D;
    // 映射至高地址
    index = (VM_START >> 30) & 511;
    early_pgtbl[index] = ppn2 | PTE_V | PTE_R | PTE_W | PTE_X | PTE_A | PTE_D; 
    
    printk("SETUP_VM is OK\n");
}


void setup_vm_final(void) {

    memset(swapper_pg_dir, 0x0, PGSIZE);
    uint64 va,pa,text_size,rodata_size,other_size;
    text_size = _srodata - _stext;
    rodata_size = _sdata - _srodata;
    other_size = PHY_SIZE - (_sdata-_stext);

    // No OpenSBI mapping required

    // mapping kernel text X|-|R|V
    va = VM_START + OPENSBI_SIZE;
    pa = PHY_START + OPENSBI_SIZE;
    create_mapping(swapper_pg_dir, va, pa, text_size, PTE_X | PTE_R | PTE_V | PTE_A | PTE_D);

    // mapping kernel rodata -|-|R|V
    va += text_size;
    pa += text_size;
    create_mapping(swapper_pg_dir, va, pa, rodata_size, PTE_R | PTE_V | PTE_A | PTE_D);
  
    // mapping other memory -|W|R|V
    va += rodata_size;
    pa += rodata_size;
    create_mapping(swapper_pg_dir, va, pa, other_size, PTE_W | PTE_R | PTE_V | PTE_A | PTE_D);

    // write satp register  
    uint64 satpVal = (((uint64)(swapper_pg_dir) - PA2VA_OFFSET) >> 12) | ((uint64)(8) << 60);
    csr_write(satp, satpVal); 

    // flush TLB
    asm volatile("sfence.vma zero, zero");

    // flush icache
    asm volatile("fence.i");
    return;
}


/* 创建多级页表映射关系 */
void create_mapping(uint64 *pgtbl, uint64 va, uint64 pa, uint64 sz, uint64 perm) {
    /*
    pgtbl 为根页表的基地址
    va, pa 为需要映射的虚拟地址、物理地址
    sz 为映射的大小
    perm 为映射的读写权限

    将给定的一段虚拟内存映射到物理内存上
    物理内存需要分页
    创建多级页表的时候可以使用 kalloc() 来获取一页作为页表目录
    可以使用 V bit 来判断页表项是否存在
    */
    uint64 page_num = 0, *table, pte, index, ppn;
    for( uint64 page_num = 0;page_num < sz;page_num += PGSIZE){
        table = pgtbl;      // 这是个虚拟地址
        index = (va >> 30) & 0x1ff; // VPN2
        pte = table[index];
        if(!(pte & PTE_V)){     // 如果当前表项的 V 位是0
                                // 我们需要新建一个页表
            // 在物理空间上新建一个页
            uint64 newpageAddr = kalloc() - PA2VA_OFFSET; // kalloc会返回一个虚拟地址
                                                            // newtable指向一个物理空间上的物理地址
            memset((void *)newpageAddr,0, PGSIZE);
            // 根据物理地址的PPN填入pte的表项
            ppn = (newpageAddr >> 12) << 10;
            pte = ppn | PTE_V;
            table[index] = pte;
        }

        table = (uint64 *)(((pte >> 10) << 12) + PA2VA_OFFSET);   // 得到虚拟地址
        index = (va >> 21) & 0x1ff; //VPN1
        pte = table[index];
        if(!(pte & PTE_V)){     // 如果当前表项的 V 位是0
                                // 我们需要新建一个页表
            // 在物理空间上新建一个页
            uint64 newpageAddr = kalloc() - PA2VA_OFFSET; // kalloc会返回一个虚拟地址
                                                            // newtable指向一个物理空间上的物理地址
            memset((void *)newpageAddr,0, PGSIZE);
            // 根据物理地址的PPN填入pte的表项
            ppn = (newpageAddr >> 12) << 10;
            pte = ppn | PTE_V;
            table[index] = pte;
        }

        table = (uint64 *)(((pte >> 10) << 12) + PA2VA_OFFSET);
        index = (va >> 12) & 0x1ff; // VPN0

        pte = (pa >> 12) << 10 | PTE_V | perm;
        table[index] = pte;

        va += PGSIZE;
        pa += PGSIZE;
    }


}
